from ftplib import *
from concurrent.futures import ThreadPoolExecutor
import pathlib
import glob, os, re

'''
1. ftp 基础功能加强;
2. downloader 和 uploader 各自实现;
3. 内容提供, 以便多线程开发; 数据管理;
4. 
'''


class MyFTP(FTP):
    def __init__(self, *list, **dict):
        super().__init__(*list, **dict)
        self.path_queue = []

    def pushd(self, path):
        before = self.pwd()
        self.cwd(path)
        self.path_queue.append(before)

    def popd(self):
        before = self.path_queue[-1]
        self.cwd(before)
        self.path_queue.pop()

    def goto_abs(self, path):
        ppath = pathlib.PurePosixPath(path)
        if ppath.root != "/":
            raise Exception(f"abs path must be start with /, input '{path}'")
        abs_path =  ppath.as_posix()

        self.cwd(path)
        abs_pwd = self.pwd()
        if abs_pwd != abs_path:
            raise Exception(f"{abs_pwd} != {abs_path}")

    def is_file(self, path):
        try:
            self.sendcmd(f"MDTM {path}")
            return True
        except:
            return False

    def mkd_p(self, path):
        ppath = pathlib.PurePosixPath(path)
        if ppath.name != "":
            self.mkd_p(ppath.parent.as_posix())
            try:
                cmd = f"NLST {path}"
                self.retrlines(cmd, lambda x: x)
            except error_perm as e:
                '''
                    this is caused by dir not exist
                '''
                if e.args[0].startswith("550 "):
                    self.mkd(path)
                else:
                    raise e
            except Exception as e:
                raise e

    def walk(self, path):
        return self._walk(path)

    def _walk(self, path):
        if self.is_file(path):
            return [["./" , [path]]]
        else:
            ret_val = []
            files = []
            lists = []
            self.pushd(path)
            self.retrlines(f"LIST", lists.append)
            for file in lists:
                base_name = re.match("(?:[\\w:-]+\\s*){7}\\w{2}:\\w{2} (.*)", file).group(1)
                if file[0] == 'd':
                    ret_val += self._walk(pathlib.PurePosixPath(path).joinpath(base_name).as_posix())
                elif file[0] == '-':
                    files.append(base_name)
                else:
                    print(f"{file } ignored")
            ret_val.append([pathlib.PurePosixPath(path).as_posix(), files])
            self.popd()
            return ret_val


    def upload(self):
        '''
            case 1: file 2 file
            case 2: dir 2 dir
        '''
        pass

class Downloader:
    def __init__(self, ftp:MyFTP):
        self.ftp = ftp

    def download_r(self, dipr, cur_path):
        if os.path.exists(cur_path) and not os.path.isdir(cur_path):
            raise Exception(f"{cur_path} is not a dir path")
        self.download_list(self.ftp.walk(dipr), os.path.abspath(cur_path))

    def download_list(self, list, dst_dir):
        for path, files in list:
            out_path = os.path.abspath(dst_dir + os.path.sep + path)
            os.makedirs(out_path, exist_ok=True)
            for file in files:
                ftp_path_obj = pathlib.PurePosixPath(path).joinpath(file)
                src_file = ftp_path_obj.as_posix()
                dst_file = os.path.join(out_path, ftp_path_obj.name)
                self._download(src_file, dst_file)

    def _download(self, src_file, dst_file):
        print(f"'{src_file}' => '{dst_file}'")
        with open(dst_file, "wb") as fd:
            self.ftp.retrbinary(f"RETR {src_file}", fd.write)


class Uploader:
    def __init__(self, ftp:MyFTP):
        self.ftp = ftp

    def upload_r(self, remote_path, *dirs):
        # if not self.ftp.is_file(remote_path):
        #     raise Exception(f"{remote_path} is not a dir path")

        '''
            walk all dirs
        '''
        upload_list = []
        for dir in dirs:
            for path, subdirs, subfiles in os.walk(dir):
                upload_list.append([dir, os.path.relpath(path, dir), subfiles])
        # print(upload_list)
        self.upload_list(upload_list, remote_path)

    def upload_list(self, list, dst_dir):
        self.ftp.pushd(dst_dir)
        for abs_dir, path, files in list:
            dst_path = pathlib.PurePath(path).as_posix()
            self.ftp.mkd_p(dst_path)
            for file in files:
                self._upload(os.path.join(abs_dir, path, file), pathlib.PurePath(path).joinpath(file).as_posix())
        self.ftp.popd()

    def _upload(self, src_file, dst_file):
        print(f"'{src_file}' => '{dst_file}'")
        with open(src_file, "rb") as fd:
            self.ftp.storbinary(f"STOR {dst_file}", fd)

def ftp():
    download_path = "/"
    store_path = r"D:\\temp"

    ftp = MyFTP()
    ftp.connect("127.0.0.1", port=2121)
    ftp.login()
    result = ftp.walk(download_path)
    ftp.close()

    thread_pool_count= 6
    with ThreadPoolExecutor(max_workers=thread_pool_count) as executor:
        def download(lists, dst):
            t = MyFTP()
            t.connect("127.0.0.1", port=2121)
            t.login()
            Downloader(t).download_list(lists, dst)
            t.close()
        rangeres = thread_pool_count
        lenres = len(result)
        step = max(1, len(result) // rangeres)
        futures = [executor.submit(download, result[i*step:min(lenres, (i+1)*step)], store_path) for i in range(0, rangeres - 1)]
        for future in futures:
            future.result()

# ftp()

def upload():
    t = MyFTP()
    t.connect("127.0.0.1", port=2121)
    t.login()
    t.mkd_p("/wtf/man/")
    Uploader(t).upload_r("/wtf/man/", r"D:\temp")
    t.close()

upload()

# print(glob.glob("*"))